import tensorflow as tf
from captcha.image import ImageCaptcha
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import random
import tensorlayer as tl

import tensorflow.contrib.slim as slim

number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
max_captcha = 4
char_set_len = 10
image_height = 60
image_width = 160


# 生成验证码的字符串
def random_captcha_text(char_set=number, captcha_size=4):
    captcha_text = []
    for i in range(captcha_size):
        c = random.choice(char_set)
        captcha_text.append(c)
    return captcha_text


# 生成验证码
def gen_captcha_text_image():
    image = ImageCaptcha()
    captcha_text = random_captcha_text()
    captcha_text = ''.join(captcha_text)
    captcha = image.generate(captcha_text)
    captcha_image = Image.open(captcha)
    captcha_image = np.array(captcha_image)
    return captcha_text, captcha_image


# 彩色转灰度图
def convert2gray(img):
    if len(img.shape) > 2:
        r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2]
        gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
        return gray
    else:
        return img


# 数字转向量
def text2vec(text):
    text_len = len(text)
    if text_len > max_captcha:
        raise ValueError('验证码最长4个字符')

    vector = np.zeros(max_captcha * char_set_len)

    def char2pos(c):
        if c == '_':
            k = 62
            return k
        k = ord(c) - 48
        if k > 9:
            k = ord(c) - 55
            if k > 35:
                k = ord(c) - 61
                if k > 61:
                    raise ValueError('No Map')
        return k

    for i, c in enumerate(text):
        idx = i * char_set_len + char2pos(c)
        vector[idx] = 1
    return vector


# 生成一个批次的数据
def get_next_batch(batch_size=128):
    batch_x = np.zeros([batch_size, image_height * image_width])  # [,60*160]
    batch_y = np.zeros([batch_size, max_captcha * char_set_len])  # [,4*10]

    def wrap_gen_captcha_text_and_image():
        while True:
            text, image = gen_captcha_text_image()
            if image.shape == (60, 160, 3):
                return text, image

    for i in range(batch_size):
        text, image = wrap_gen_captcha_text_and_image()
        image = convert2gray(image)  # 转成灰度图

        batch_x[i, :] = image.flatten() / 255  # 返回一个折叠成 一维 的数组
        batch_y[i, :] = text2vec(text)

    return batch_x, batch_y


X = tf.placeholder(tf.float32, [None, image_height * image_width])
Y = tf.placeholder(tf.float32, [None, max_captcha * char_set_len])
keep_prob = tf.placeholder(tf.float32)  # dropout


# keep_prob=0.75

def crack_captcha_cnn():
    # 建立模型
    x = tf.reshape(X, shape=[-1, image_height, image_width, 1])
    # print(x)
    # 将占位符 转换为 按照图片给的新样式
    # slim.conv2d的stride默认为1，slim.max_pool2d的stride默认与f保持一致, padding默认都为SAME
    x = tf.reshape(X, shape=[-1, image_height, image_width, 1])
    with slim.arg_scope([slim.conv2d, slim.fully_connected], activation_fn=tf.nn.relu,
                        weights_initializer=tf.truncated_normal_initializer(stddev=0.1),
                        weights_regularizer=slim.l2_regularizer(0.0005)):
        net = slim.conv2d(x, 32, 3, scope="conv1")  # 32为通道，3为f即[3, 3]，stride默认为1，padding默认为SAME
        net = slim.max_pool2d(net, 2, scope="pool1")  # 2即f为[2,2],stride默认与f保持一致
        net = slim.conv2d(net, 64, 1, scope="conv2")
        net = slim.max_pool2d(net, 2, scope="pool2")
        net = slim.conv2d(net, 64, 1, scope="conv3")
        net = slim.max_pool2d(net, 2, scope="pool3")
        net = slim.flatten(net, scope='flat_1')
        net = slim.dropout(net, keep_prob=keep_prob, scope='drop1')
        net = slim.fully_connected(net, 1024, scope="f1")
        net = slim.fully_connected(net, 512, scope="f2")
        z = slim.fully_connected(net, max_captcha * char_set_len, activation_fn=None, scope="f3")
        p = tf.nn.sigmoid(z, name="prediction")
    return z, p


def train():
    # sess = tf.InteractiveSession()
    # 定义损失函数
    # network = net()  #
    # z = network.outputs
    # p = tf.nn.sigmoid(z, name="prediction")
    z, p = crack_captcha_cnn()

    # sigmoid多任务模式，一个标签中有多个目标
    # print(z)
    # print(Y)
    loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=z, labels=Y))
    # loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=z, labels=Y))
    # 定义优化器
    # train_param = network.all_params
    optimizer = tf.train.AdamOptimizer(learning_rate=0.001).minimize(loss)
    predict = tf.reshape(p, [-1, max_captcha, char_set_len])
    max_idx_p = tf.argmax(predict, 2)
    max_idx_l = tf.argmax(tf.reshape(Y, [-1, max_captcha, char_set_len]), 2)
    correct_pred = tf.equal(max_idx_p, max_idx_l)
    accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        step = 0
        while True:
            batch_x, batch_y = get_next_batch(64)
            # batch_y = np.asarray(batch_y, dtype=np.float32)
            # batch_x = np.asarray(batch_x, dtype=np.float32)

            # print(batch_x.shape)
            # print(batch_y.shape)
            _, loss_ = sess.run([optimizer, loss], feed_dict={X: batch_x, Y: batch_y, keep_prob: 0.75})
            print(step, loss_)

            # 每100 step计算一次准确率
            if step % 100 == 0:
                batch_x_test, batch_y_test = get_next_batch(100)
                acc = sess.run(accuracy, feed_dict={X: batch_x_test, Y: batch_y_test, keep_prob: 1.})
                print(step, acc)
                # 如果准确率大于50%,保存模型,完成训练
                if acc > 0.9:
                    saver.save(sess, "models\\crack_capcha.model", global_step=step)
                    break
            step = step + 1


if __name__ == '__main__':
    train()
